import random
from shutil import copy
from pathlib import Path

def ratio_split(data_folder, points_folder, ratio):
    # 获取数据目录
    data_dir = Path(data_folder)
    # 动态获取类别目录
    categories = [d.name for d in data_dir.iterdir() if d.is_dir()]
    # 创建目标目录
    train_dir = Path(points_folder)
    train_dir = train_dir / 'train'
    val_dir = Path(points_folder)
    val_dir = val_dir / 'val'
    train_dir.mkdir(parents=True, exist_ok=True)
    val_dir.mkdir(parents=True, exist_ok=True)
    
    # 遍历每个类别
    for category in categories:
        # 获取类别目录
        category_dir = data_dir / category
        
        # 获取类别目录下的所有文件
        files = list(category_dir.glob('*'))
        
        # 打乱文件顺序
        random.shuffle(files)
        
        # 计算拆分点
        split_index = int(len(files) * ratio)
        
        # 拆分文件
        train_files = files[:split_index]
        val_files = files[split_index:]
        
        # 创建类别子目录
        train_category_dir = train_dir / category
        val_category_dir = val_dir / category
        train_category_dir.mkdir(parents=True, exist_ok=True)
        val_category_dir.mkdir(parents=True, exist_ok=True)
        
        # 移动文件
        for file in train_files:
            copy(file, train_category_dir)
        for file in val_files:
            copy(file, val_category_dir)

# 示例调用
if __name__ == '__main__':
    ratio_split('../data', 0.8)